from typing import List
import streamlit as st
import google.generativeai as genai
from core.config import AgentConfig
from core.gemini import GeminiAgent
import google.generativeai.types.content_types as content_types

@st.cache_resource
def create_chat_model(gemini:dict[str,str]):
    api_key = gemini["api_key"]
    genai.configure(api_key=api_key,
                    transport=gemini["transport"])
    model = genai.GenerativeModel(model_name=gemini["model_name"],
        system_instruction=gemini["system_instruction"])
    return model

def call_exec(script:str):
    """
    调用python exec函数
    注意：生成的脚本中不允许包含 `if __name__ == '__main__': pass` 
    这样的启动代码，否则不执行。
    """
    global_vars={
        "exec_result": None
    }
    try:
        compiled_code = compile(script,"","exec")
        exec(compiled_code,global_vars)
    except Exception as e:
        raise e
    
    return global_vars

def chat_container(config:AgentConfig):
    model = create_chat_model(config.gemini)
    gemini_agent = GeminiAgent(model=model)

    # 聊天记录
    if "messages" not in st.session_state:
        st.session_state["messages"] = [
            {"role": "assistant", 
            "content": config.streamlit['welcome_word']}
        ]
    #遍历聊天记录    
    for index,msg in enumerate(st.session_state.messages):
        st.chat_message(msg["role"]).write(msg["content"])
        #按钮右边布局
        col1,col2= st.columns([6,1])
        if index and msg['role']=='assistant':
            if col2.button("执行脚本",key='bnt{}'.format(index),type='primary'):
                script = gemini_agent.extract_python_script(msg["content"])
                if script:
                    try:
                        ret = call_exec(script)
                        msg["exec_result"]=ret["exec_result"]
                        # st.chat_message("ai").write(ret["exec_result"])
                    except Exception as e:
                        col1.error(f"脚本执行错误：{e.args[0]}")
                        st.stop()
            #如果有执行结果则显示
            if msg["exec_result"]:
                st.chat_message("ai").write(msg["exec_result"])
    if prompt := st.chat_input("请描述您的需求。",max_chars=1000):
        #添加对话到历史记录
        st.session_state.messages.append({"role": "user", "content": prompt})
        st.chat_message("user").write(prompt)

        #响应消息
        with st.chat_message("assistant"):
            with st.spinner("正在思考..."):
                response = gemini_agent.send_message(prompt=prompt)
                answer = (part.text for part in response)
                full_answer =  st.write_stream(answer)
                st.session_state.messages.append({"role": "assistant", "content": full_answer,"exec_result":None})
            if full_answer is not None:
                st.rerun()